﻿#include "ipv4_frag.h"
#include "mblock.h"
#include "pktbuf.h"
#include "timer.h"

static ip_frag_t frag_entry[IP_FRAGS_MAX_NR]; // 用于存储分片数据包表项
static mblock_t frag_mblock;                  // 分片数据包表项内存管理器
static nlist_t frag_list;                     // 存放分片数据包表项的链表
static net_timer_t frag_timer;                // 分片定时器

#if DBG_DISP_ENABLED(DBG_IPV4_LEVEL)
static void display_ip_frag()
{
    plat_printf("-------------ipfrags:----------------\n");

    nlist_node_t *fp;
    int f_index = 0;
    nlist_for_each(fp, &frag_list)
    {
        ip_frag_t *frag = nlist_entry(ip_frag_t, node, fp);
        plat_printf("[frag_entry:%d]:\n", f_index++);
        dbg_dump_ipaddr("ip:", &frag->ip);
        plat_printf("id:%d, tmo:%d\n", frag->id, frag->tmo);

        nlist_node_t *pp;
        int p_index = 0;
        nlist_for_each(pp, &frag->buf_list)
        {
            pktbuf_t *buf = nlist_entry(pktbuf_t, node, pp);
            ipv4_pkt_t *pkt = (ipv4_pkt_t *)pktbuf_data(buf);
            plat_printf("B%d:[%d - %d]  ", p_index++, ipfrag_get_start(pkt), ipfrag_get_end(pkt));
        }
        plat_printf("\n");
    }

    plat_printf("-------------------------------------\n");
}
#else
#define display_ip_frag()
#endif

/**
 * @brief 获取分片数据包在整个数据包中的起始地址
 */
int ipfrag_get_start(ipv4_pkt_t *pkt)
{
    return pkt->hdr.offset * 8;
}

/**
 * @brief 获取分片数据包的起始长度
 */
int ipfrag_get_data(ipv4_pkt_t *pkt)
{
    return pkt->hdr.total_len - ipv4_hdr_size(pkt);
}

/**
 * @brief 获取分片数据包在整个数据包中的结束地址
 */
int ipfrag_get_end(ipv4_pkt_t *pkt)
{
    return ipfrag_get_start(pkt) + ipfrag_get_data(pkt) - 1;
}

/**
 * @brief 释放分片表项所有被存储的数据包
 */
static void ipfrag_free_buflist(ip_frag_t *frag)
{
    nlist_node_t *node;
    while ((node = nlist_remove_first(&frag->buf_list)))
    {
        pktbuf_t *buf = nlist_entry(pktbuf_t, node, node);
        pktbuf_free(buf);
    }
}

/**
 * @brief 根据ip地址和数据包id来寻找对应的分片表项
 * @param ip 分片数据包的源ip地址
 * @param id 分片数据包的id
 */
static ip_frag_t *ipfrag_find(const ipaddr_t *ip, uint16_t id)
{
    nlist_node_t *p = NULL;
    nlist_for_each(p, &frag_list)
    {
        ip_frag_t *frag = nlist_entry(ip_frag_t, node, p);
        if (ipaddr_is_equal(ip, &frag->ip) && id == frag->id)
        {
            // 查找完以后，把分片表项放在前面，减少后续遍历次数
            nlist_node_t *node = nlist_remove(&frag_list, &frag->node);
            nlist_insert_first(&frag_list, node);
            return frag;
        }
    }

    return NULL;
}

/**
 * @brief 向分片表项中添加表项
 * @param frag 已经分配好的表项
 * @param ip ip地址
 * @param id 数据包id
 */
static void ipfrag_add(ip_frag_t *frag, ipaddr_t *ip, uint16_t id)
{
    ipaddr_copy(&frag->ip, ip);
    frag->id = id;
    nlist_init(&frag->buf_list);
    nlist_node_init(&frag->node);
    frag->tmo = IP_FRAG_TMO / IP_FRAGS_SCAN_PERIOD;

    nlist_insert_first(&frag_list, &frag->node);
}

/**
 * @brief 判断分片表项中的分片数据包是否已经全部到达
 */
static bool ipfrag_is_all_arrived(ip_frag_t *frag)
{
    uint16_t offset = 0;
    nlist_node_t *p = NULL;
    ipv4_pkt_t *pkt = NULL;
    nlist_for_each(p, &frag->buf_list)
    {
        pktbuf_t *buf = nlist_entry(pktbuf_t, node, p);
        pkt = (ipv4_pkt_t *)pktbuf_data(buf);

        if (offset != ipfrag_get_start(pkt))
        {
            return false;
        }

        offset += ipfrag_get_data(pkt);
    }

    return pkt ? !pkt->hdr.more : false;
}

/**
 * @brief 将分片表项中的数据包合并为一个数据包
 */
static pktbuf_t *ipfrag_join(ip_frag_t *frag)
{
    nlist_node_t *node;
    pktbuf_t *target = NULL;

    while ((node = nlist_remove_first(&frag->buf_list)))
    {
        pktbuf_t *curr_buf = nlist_entry(pktbuf_t, node, node);
        ipv4_pkt_t *curr_pkt = (ipv4_pkt_t *)pktbuf_data(curr_buf);
        if (!target)
        {
            target = curr_buf;
            continue;
        }

        // 先移除包头再合并
        net_err_t err = pktbuf_remove_header(curr_buf, ipv4_hdr_size(curr_pkt));
        if (err < NET_ERR_OK)
        {
            DBG_ERROR(DBG_IPV4_LEVEL, "remove header:%d failed!", ipv4_hdr_size(curr_pkt));
            pktbuf_free(curr_buf);
            goto free_and_return;
        }

        err = pktbuf_join(target, curr_buf);
        if (err < NET_ERR_OK)
        {
            DBG_ERROR(DBG_IPV4_LEVEL, "pktbuf join failed!");
            pktbuf_free(curr_buf);
            goto free_and_return;
        }
    }

    DBG_INFO(DBG_IPV4_LEVEL, "the pkt in frag_entry:%d is all arrived!", frag->id);
    ipfrag_free(frag);
    return target;

free_and_return:
    ipfrag_free(frag);
    pktbuf_free(target);
    return NULL;
}

/**
 * @brief 按顺序在分片表项中插入数据包
 * @param frag 被插入的表项
 * @param buf 收到的数据包
 * @param pkt 收到的ipv4类型的数据包
 */
static net_err_t ipfrag_insert(ip_frag_t *frag, pktbuf_t *buf, ipv4_pkt_t *pkt)
{
    // 先判断表项中数据包的数量，避免过度消耗内存
    if (nlist_count(&frag->buf_list) > IPFRAGS_MAX_BUF_NR)
    {
        DBG_WARNING(DBG_IPV4_LEVEL, "frag buf list is full!");
        return NET_ERR_FULL;
    }

    // 遍历找到正确的位置并插入
    nlist_node_t *p;
    nlist_for_each(p, &frag->buf_list)
    {
        pktbuf_t *curr_buf = nlist_entry(pktbuf_t, node, p);
        ipv4_pkt_t *curr_pkt = (ipv4_pkt_t *)pktbuf_data(curr_buf);
        uint16_t curr_start = ipfrag_get_start(curr_pkt);

        if (curr_start == ipfrag_get_start(pkt))
        {
            DBG_WARNING(DBG_IPV4_LEVEL, "the ipv4 frag is existed! pkt id:%d\n", pkt->hdr.id);
            return NET_ERR_EXISTED;
        }
        else if (curr_start > ipfrag_get_end(pkt))
        {
            nlist_node_t *pre = nlist_node_pre(p);
            if (pre)
            {
                nlist_insert_after(&frag->buf_list, pre, &buf->node);
            }
            else
            {
                nlist_insert_first(&frag->buf_list, &buf->node);
            }

            return NET_ERR_OK;
        }
    }

    nlist_insert_last(&frag->buf_list, &buf->node);
    return NET_ERR_OK;
}

/**
 * @brief 分配分片数据包表项
 * @param force 空间不足时，是否强制分配
 * @return 分配的表项
 */
ip_frag_t *ipfrag_alloc(bool force)
{
    ip_frag_t *frag = mblock_alloc(&frag_mblock, DELIVER_NO_WAITTIME);
    if (!frag)
    {
        DBG_INFO(DBG_IPV4_LEVEL, "ipfrag alloc failed and will force alloc!");
        nlist_node_t *node = nlist_remove_last(&frag_list);
        if (node)
        {
            frag = nlist_entry(ip_frag_t, node, node);
            ipfrag_free_buflist(frag);
        }
    }

    return frag;
}

/**
 * @brief 释放分片表项
 * @param frag 要释放的分片表项
 */
void ipfrag_free(ip_frag_t *frag)
{
    ipfrag_free_buflist(frag);
    nlist_remove(&frag_list, &frag->node);
    mblock_free(&frag_mblock, frag);
}

/**
 * @brief 分片定时器超时后调用的处理函数，清楚超时的分片表项
 */
static void ipfrag_tmo(net_timer_t *timer, void *arg)
{
    nlist_node_t *curr = NULL, *next = NULL;

    DBG_INFO(DBG_IPV4_LEVEL, "scan frag!");

    for (curr = nlist_first(&frag_list); curr; curr = curr->next)
    {
        next = nlist_node_next(curr);

        ip_frag_t *curr_frag = nlist_entry(ip_frag_t, node, curr);
        if (--curr_frag->tmo <= 0)
        {
            DBG_INFO(DBG_IPV4_LEVEL, "the id:%d frag is time out!", curr_frag->id);
            ipfrag_free(curr_frag);
        }
    }
}

/**
 * @brief 处理收到的分片数据包
 * @param netif 收到分片数据包的网络接口
 * @param buf 收到的分片数据包
 * @param src_ip 收到的分片数据包的源ip地址
 * @param dst_ip 收到的分片数据包的目的ip地址
 */
net_err_t ipfrag_in(netif_t *netif, pktbuf_t *buf, ipaddr_t *src_ip, ipaddr_t *dst_ip)
{
    // 进入到这个函数里，应该已经设置过包的连续了
    ipv4_pkt_t *pkt = (ipv4_pkt_t *)pktbuf_data(buf);
    ip_frag_t *frag = ipfrag_find(src_ip, pkt->hdr.id);
    if (!frag)
    {
        frag = ipfrag_alloc(true);
        if (!frag)
        {
            DBG_ERROR(DBG_IPV4_LEVEL, "ip frag in failed!");
            pktbuf_free(buf);
            return NET_ERR_OK;
        }
        ipfrag_add(frag, src_ip, pkt->hdr.id);
    }
    net_err_t err = ipfrag_insert(frag, buf, pkt);
    if (err < NET_ERR_OK)
    {
        DBG_WARNING(DBG_IPV4_LEVEL, "frag insert buf failed!");
        pktbuf_free(buf);
        ipfrag_free(frag);
        return NET_ERR_OK;
    }

    display_ip_frag();

    if (ipfrag_is_all_arrived(frag))
    {
        pktbuf_t *target_buf = ipfrag_join(frag);
        if (!target_buf)
        {
            DBG_ERROR(DBG_IPV4_LEVEL, "ipfrag join failed!");
            display_ip_frag();
        }
        else
        {
            DBG_INFO(DBG_IPV4_LEVEL, "the frag is all arived! the join buf total size:%d!", target_buf->total_size);
            err = ipv4_normal_in(netif, target_buf, dst_ip, src_ip);
            if (err < NET_ERR_OK)
            {
                DBG_WARNING(DBG_IPV4_LEVEL, "frag but to ipv4 normal in failed!");
                pktbuf_free(target_buf);
            }
        }

        ipfrag_free(frag);
    }

    return NET_ERR_OK;
}

/**
 * @brief 数据包分片重组相关初始化
 */
net_err_t ipfrag_init()
{
    nlist_init(&frag_list);
    net_err_t err = mblock_init(&frag_mblock, frag_entry, sizeof(ip_frag_t), IP_FRAGS_MAX_NR, NLOCKER_NONE);
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_IPV4_LEVEL, "ipfrag init failed!");
        return err;
    }

    // 添加分片定时器
    err = timer_add(&frag_timer, "frag_timer", ipfrag_tmo, NULL, IP_FRAGS_SCAN_PERIOD * 1000, true);
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_IPV4_LEVEL, "add ipfrag timer failed!");
        return err;
    }

    return NET_ERR_OK;
}